import numpy as np
import matplotlib.pyplot as plt
import sys
import platform

# Set font based on operating system
if platform.system() == 'Windows':
    plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows common Chinese font
elif platform.system() == 'Darwin':  # macOS
    plt.rcParams['font.sans-serif'] = ['STHeiti']
else:  # Linux or other
    plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']  # Example font, may vary based on system

# Set minus sign to display correctly
plt.rcParams['axes.unicode_minus'] = False

# 生成带噪声的线性数据
def generate_data():
    np.random.seed(42)
    X = 2 * np.random.rand(100, 1)
    y = 4 + 3 * X + np.random.randn(100, 1)  # y = 4 + 3X + 噪声
    return X, y

# 显示用法说明和错误消息
def show_instructions():
    print("用法: linear_regression.exe <学习率> <迭代次数>")
    print("示例: linear_regression.exe 0.1 20")
    print("其中 <学习率> 是一个小数值，例如 0.1")
    print("      <迭代次数> 是一个正整数，例如 20")

# 计算均方误差（MSE）
def compute_mse(y, y_pred):
    return np.mean((y - y_pred) ** 2)

# 梯度下降函数用于线性回归，并可视化收敛情况
def gradient_descent(X, y, theta_0, theta_1, learning_rate, iterations):
    m = len(X)
    mse_history = []  # 记录每次迭代的 MSE 值

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 10))
    plt.ion()  # 启用交互模式，实时更新图形

    for i in range(1, iterations + 1):
        # 计算预测值
        y_pred = theta_0 + theta_1 * X
        
        # 计算均方误差
        mse = compute_mse(y, y_pred)
        mse_history.append(mse)
        
        # 计算梯度
        d_theta_0 = (1/m) * np.sum(y_pred - y)
        d_theta_1 = (1/m) * np.sum((y_pred - y) * X)
        
        # 更新参数
        theta_0 -= learning_rate * d_theta_0
        theta_1 -= learning_rate * d_theta_1

        # 可视化当前状态
        ax1.cla()  # 清除上一帧的绘图
        ax1.scatter(X, y, color='blue', alpha=0.5, label="数据点")
        ax1.plot(X, theta_0 + theta_1 * X, color='red', label=f"y = {theta_0:.2f} + {theta_1:.2f}x")
        ax1.set_title(f"线性回归 - 第 {i} 次迭代")
        ax1.set_xlabel("X")
        ax1.set_ylabel("y")
        ax1.legend()

        # 在图形上显示附加信息
        ax1.text(0.05, 0.95, f"迭代次数: {i}", transform=ax1.transAxes, fontsize=10, verticalalignment='top')
        ax1.text(0.05, 0.90, f"截距 (Theta_0): {theta_0:.4f}", transform=ax1.transAxes, fontsize=10, verticalalignment='top')
        ax1.text(0.05, 0.85, f"斜率 (Theta_1): {theta_1:.4f}", transform=ax1.transAxes, fontsize=10, verticalalignment='top')
        ax1.text(0.05, 0.80, f"均方误差 (MSE): {mse:.4f}", transform=ax1.transAxes, fontsize=10, verticalalignment='top')

        # 第二个子图绘制每次迭代的 MSE 变化情况
        ax2.cla()
        ax2.plot(mse_history, label="迭代过程中的 MSE", color='purple')
        ax2.set_title("MSE 收敛情况")
        ax2.set_xlabel("迭代次数")
        ax2.set_ylabel("均方误差")
        ax2.legend()

        # 检查是否收敛（MSE 变化很小）
        if i > 1 and abs(mse_history[-2] - mse_history[-1]) < 1e-4:
            ax2.text(0.5, 0.5, "已收敛", transform=ax2.transAxes, fontsize=14, color="green", ha="center")
        
        # 显示高学习率导致的“不稳定”警告
        if i > iterations * 0.8 and mse > min(mse_history):
            ax2.text(0.5, 0.7, "警告: 学习率过高导致的不稳定", transform=ax2.transAxes, fontsize=12, color="red", ha="center")

        plt.pause(0.5)  # 暂停以创建动画效果
    
    plt.ioff()  # 关闭交互模式
    plt.show()
    return theta_0, theta_1

def main():
    # 验证命令行参数
    if len(sys.argv) != 3:
        print("错误: 参数数量不正确。")
        show_instructions()
        sys.exit(1)
    
    try:
        # 解析命令行参数
        learning_rate = float(sys.argv[1])
        iterations = int(sys.argv[2])
    except ValueError:
        print("错误: 请确保 <学习率> 是小数，<迭代次数> 是整数。")
        show_instructions()
        sys.exit(1)

    # 生成数据并初始化参数
    X, y = generate_data()
    theta_0 = 0  # 截距
    theta_1 = 0  # 斜率

    # 执行梯度下降并可视化过程
    print(f"正在运行线性回归，学习率 = {learning_rate}，迭代次数 = {iterations}...")
    theta_0, theta_1 = gradient_descent(X, y, theta_0, theta_1, learning_rate, iterations)
    print(f"最终模型: y = {theta_0:.2f} + {theta_1:.2f}x")

if __name__ == "__main__":
    main()